package cn.jishuba.spring.rag;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.CommandLineRunner;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.WebApplicationType;
import org.springframework.boot.autoconfigure.SpringBootApplication;

import java.util.List;

@SpringBootApplication
public class SpringRagDemoApplication implements CommandLineRunner {

    public static void main(String[] args) {
        SpringApplication app = new SpringApplication(SpringRagDemoApplication.class);
        app.setWebApplicationType(WebApplicationType.NONE);
        app.run(args);
    }

    @Override
    public void run(String... args) throws Exception {
        // 读取PDF文档
        PagePdfDocumentReader reader = new PagePdfDocumentReader("classpath:2025-Microsoft-AI-in-Education-Report.pdf");
        List<Document> documents = reader.read();

        // 初始化OllamaApi
        OllamaApi ollamaApi = OllamaApi.builder().build();

        // 使用OllamaEmbeddingModel和SimpleVectorStore来处理PDF文档
        EmbeddingModel embeddingModel = OllamaEmbeddingModel.builder()
                .ollamaApi(ollamaApi)
                .defaultOptions(OllamaOptions.builder()
                        .model("qwen2.5:latest")
                        .build())
                .build();
        VectorStore vectorStore = SimpleVectorStore.builder(embeddingModel)
                .build();

        // 将文档添加到向量存储中
        vectorStore.add(documents);

        // 构建 Ollama ChatModel
        ChatModel chatModel = OllamaChatModel.builder()
                .ollamaApi(ollamaApi)
                .defaultOptions(OllamaOptions.builder()
                        .model("qwen2.5:latest")
                        .build())
                .build();

        // 使用 ChatClient 和 QuestionAnswerAdvisor 来查询向量存储
        var result = ChatClient.builder(chatModel)
                .build()
                .prompt()
                .advisors(new QuestionAnswerAdvisor(vectorStore)) // 使用向量存储作为顾问
                .user("请概括文档中的内容")
                .call()
                .content();

        System.out.println(result);
    }
}
